from termcolor import colored
from datetime import datetime
import json
import sys
from pydoc import pipepager

author = {
	"name":"gl4ssesbo1",
	"twitter":"https://twitter.com/gl4ssesbo1",
	"github":"https://github.com/gl4ssesbo1",
	"blog":"https://www.pepperclipp.com/"
}

needs_creds = True

variables = {
	"SERVICE": {
		"value": "ssm",
		"required": "true",
		"description":"The service that will be used to run the module. It cannot be changed."
	},
	"INSTANCE-IDS": {
		"value": "",
		"required": "false",
		"description":"The ID of the instance or several Instance IDs split by comma."
	},
	"INSTANCE-FILE": {
		"value": "",
		"required": "false",
		"description":"A file with all of the Instance IDs that you want to send command to."
	},
	"DOCUMENT-NAME": {
		"value": "",
		"required": "true",
		"description":"The name of the SSM document to define the parameters and plugin settings for the session. For example, AWS-RunShellScript, or AWS-RunPowerShellScript. You can call the GetDocument API to verify the document exists before attempting to start a session. If no document name is provided, a shell to the instance is launched by default."
	},
	"COMMAND-FILE": {
		"value": "",
		"required": "false",
		"description":"The path of the commands to execute. I you set this, it will update the Parameter File with all the commands. Alternatively, you can set the commands on Parameter File. eg "
	},
	"DOCUMENT-VERSION": {
		"value": "",
		"required": "false",
		"description":"The version of the document provided. You can specify $DEFAULT, $LATEST, or a specific version number."
	},
	"PARAMETER-FILE": {
		"value": "",
		"required": "false",
		"description":"The name of the SSM document to define the parameters and plugin settings for the session. For example, SSM-SessionManagerRunShell . You can call the GetDocument API to verify the document exists before attempting to start a session. If no document name is provided, a shell to the instance is launched by default."
	},
	"COMMENT": {
		"value": "",
		"required": "false",
		"description":"You can leave this empty, if you don't want to run a command. Even if you run, this is not necessary, but adding a convincing comment will make it harder to detect."
	},
	"TIMEOUT-SECONDS": {
		"value": "",
		"required": "false",
		"description":"The time to wait for a command to run. If this time is reached and the command has not already started running, it will not run."
	},
	"OUTPUT-S3-BUCKET-NAME": {
		"value": "",
		"required": "false",
		"description":"The name of the S3 bucket where command execution responses should be stored."
	},
	"OUTPUT-S3-BUCKET-REGION": {
		"value": "",
		"required": "false",
		"description":"The region of the S3 bucket where command execution responses should be stored."
	},
	"OUTPUT-S3-BUCKET-PREFIX": {
		"value": "",
		"required": "false",
		"description":"The directory structure within the S3 bucket where the responses should be stored."
	}
}
description = "SendCommands to some InstanceIDs provided, split by comma, or a file of InstanceIDs. The max number of IDs is 50, so if you profide more, it will ask if you want to send the request several times." \
			  "Also shows the output of the command if you want to." \
			  "The APIs used are ssm:SendCommand and ssm:GetCommandInvocation."

aws_command = """
aws ssm send-command --document-name <Document Name> --instance-ids <instance ID> --profile <profile>
aws ssm get-command-invocation --command-id <command ID> --instance-id <instance ID> --profile <profile>
"""

colors = [
	"not-used",
	"red",
	"blue",
	"yellow",
	"green",
	"magenta",
	"cyan",
	"white"
]

output = ""

def run_command(command, profile, workspace, filename, document):
	print()

def list_dictionary(d, n_tab):
	global output
	if isinstance(d, list):
		n_tab += 1
		for i in d:
			if not isinstance(i, list) and not isinstance(i, dict):
				output += ("{}{}\n".format("\t" * n_tab, colored(i, colors[n_tab])))
			else:
				list_dictionary(i, n_tab)
	elif isinstance(d, dict):
		n_tab+=1
		for key, value in d.items():
			if not isinstance(value, dict) and not isinstance(value, list):
				output += ("{}{}: {}\n".format("\t"*n_tab, colored(key, colors[n_tab], attrs=['bold']) , colored(value, colors[n_tab+1])))
			else:
				output += ("{}{}:\n".format("\t"*n_tab, colored(key, colors[n_tab], attrs=['bold'])))
				list_dictionary(value, n_tab)

def exploit(profile, workspace):
	try:
		now = datetime.now()
		dt_string = now.strftime("%d_%m_%Y_%H_%M_%S")
		file = "{}_ssm_start_session".format(dt_string)
		filename = "./workspaces/{}/{}".format(workspace, file)

		global output
		n_tab = 0

		try:
			instance_ids = variables['INSTANCE-IDS']['value']
			instance_file = variables['INSTANCE-FILE']["value"]
			documentname = variables['DOCUMENT-NAME']['value']
			documentversion = variables['DOCUMENT-VERSION']['value']
			parameters = variables['PARAMETER-FILE']['value']
			command_file = variables['COMMAND-FILE']['value']
			timeout = variables['TIMEOUT-SECONDS']['value']
			outputs3name = variables['OUTPUT-S3-BUCKET-NAME']['value']
			outputs3region = variables['OUTPUT-S3-BUCKET-REGION']['value']
			outputs3prefix = variables['OUTPUT-S3-BUCKET-PREFIX']['value']

			params = {}
			args = {}
			if not instance_ids == "" and not instance_file == "":
				print(colored(
					"[*] Either enter INSTANCE-IDs Separated by commas, or the INSTANCE-FILE, with all the Instance IDs. Not bot, not neither.",
					"red"))

			elif instance_ids == "" and instance_file == "":
				print(colored(
					"[*] Either enter INSTANCE-IDs Separated by commas, or the INSTANCE-FILE, with all the Instance IDs. Not bot, not neither.",
					"red"))
			else:
				command_array = []
				instances = []
				if not instance_ids == "":
					instances = instance_ids.split(",")
				elif not instance_file == "":
					with open(instance_file, 'r') as outputfile:
						for line in outputfile.readline():
							instances.append(line)
					outputfile.close()

				if not documentname == "":
					args['DocumentName'] = documentname

				if not documentversion == "":
					args['DocumentVersion'] = documentversion

				if not timeout == "":
					args['TimeoutSeconds'] = int(timeout)

				if not outputs3name == "":
					args['OutputS3BucketName'] = documentversion

				if not command_file == "":
					cfile = open(command_file, 'r')
					for line in cfile.readlines():
						command_array.append(line.strip())

				if not outputs3region == "":
					args['OutputS3KeyRegion'] = outputs3region

				if not outputs3prefix == "":
					args['OutputS3KeyPrefix'] = outputs3prefix

				if not parameters == "":
					with open(parameters, 'r') as pfile:
						params = json.load(pfile)
						if 'commands' in params:
							if len(params['commands']) > 0 and len(command_array) > 0:
								yn = input(
									colored("You have provided both a command file and commands inside parameter files. Which one will you use?\n1) Command File\n2) Parameter File\n>>> ", "yellow")
								)
								while int(yn.strip()) != 1 and int(yn.strip()) != 2:
									yn = input(
										colored(
											"Enter either 1 or 2\n1) Command File\n2) Parameter File\n>>> ",
											"yellow")
									)
								if int(yn.strip()) == 1:
									params['commands'] = command_array
									print(
										colored("[*] Commands in Command File Chosen!", "green")
									)

								elif int(yn.strip()) == 2:
									print(
										colored("[*] Commands in Parameter File Chosen!", "green")
									)
						else:
							if len(command_array) > 0:
								params['commands'] = command_array

					args['Parameters'] = params
				elif parameters == "" and len(command_array) > 0:
					args['Parameters'] = {
							"commands":command_array
						}

				outputfile = open(filename, 'a')
				
				yn = 'N'
				if len(instances) > 50:
					index = 0
					yn = input(
						"There are more then 50 instances, which is the limit. The script will send the command more than once, until all the instances recieve the command. Do you agree? [y/N] ")

					if yn == 'y' or yn == 'Y':
						allinst = []
						while index + 49 < len(instances):
							theinst = instances[index:index + 49]
							index += 49
							args['InstanceIds'] = theinst

							response = profile.send_command(
								**args
							)
							allinst.append(response)
						theinst = instances[index:len(instances)]
						args['InstanceIds'] = theinst
						response = profile.send_command(
							**args
						)
						allinst.append(response['Command'])

						json_data = allinst
						outputfile.write(json.dumps(json_data, indent=4, default=str))
						#outputfile.close()
						title_name = 'CommandId'

						output += colored("---------------------------------\n", "yellow", attrs=['bold'])
						for data in json_data:
							output += colored("{}: {}\n".format(title_name, data[title_name]), "yellow", attrs=['bold'])
							output += colored("---------------------------------\n", "yellow", attrs=['bold'])
							list_dictionary(data, n_tab)
							output += "\n"
							output += colored("---------------------------------\n", "yellow", attrs=['bold'])
						pipepager(output, "less -R")
						output = ""

				else:
					try:
						args['InstanceIds'] = instances
						response = profile.send_command(
							**args
						)

						json_data = response['Command']
						outputfile.write(json.dumps(json_data, indent=4, default=str))
						#outputfile.close()
						title_name = 'CommandId'

						output += colored("---------------------------------\n", "yellow", attrs=['bold'])
						output += colored("{}: {}\n".format(title_name, json_data[title_name]), "yellow",
										  attrs=['bold'])
						output += colored("---------------------------------\n", "yellow", attrs=['bold'])
						list_dictionary(json_data, n_tab)
						output += colored("---------------------------------\n", "yellow", attrs=['bold'])
						print(output)
						output = ""

						command_id = response['Command']['CommandId']
						yn = input(colored("Do you want to also get the response of the command? [Y/n] ", "yellow"))
						if yn == 'y' or 'yn' == 'Y':
							try:
								answers = []
								for instance in instances:
									response = profile.get_command_invocation(
										CommandId=command_id,
										InstanceId=instance
									)
									del response['ResponseMetadata']
									answers.append(response)
									outputfile.write(json.dumps(answers, indent=4, default=str))

								json_data = answers
								title_name = 'InstanceId'

								output += colored("---------------------------------\n", "yellow", attrs=['bold'])
								for data in json_data:
									output += colored("{}: {}\n".format(title_name, data[title_name]), "yellow",
													  attrs=['bold'])
									output += colored("---------------------------------\n", "yellow", attrs=['bold'])
									list_dictionary(data, n_tab)
									output += "\n"
									output += colored("---------------------------------\n", "yellow", attrs=['bold'])
								pipepager(output, "less -R")
								output = ""
							except:
								e = sys.exc_info()[1]
								print(colored("[*] {}".format(e), "red"))
					except:
						e = sys.exc_info()[1]
						print(colored("[*] {}".format(e), "red"))

				print(colored("[*] Command Response dumped on file '{}'.".format(filename), "green"))
				outputfile.close()


		except profile.exceptions.InvalidDocument:
			print(colored("[*] The document provided is not valid or does not exist.", "red"))

		except profile.exceptions.TargetNotConnected:
			print(colored(
				"[*] You cannot connect to the target. Either you have an internet problem, or the target might be shutdown/removed.",
				"red"))

		except:
			e = sys.exc_info()[1]
			print(colored("[*] {}".format(e), "red"))

	except:
		e = sys.exc_info()[1]
		print(colored("[*] {}".format(e), "red"))